#include "dtacore.h"
#include "libdft_api.h"
#include "tagmap.h"
#include "dta.rop.check.h"

namespace dta
{
	long	set_ins_inspect( size_t xed_icls, void * inspector )
	{
		typedef void (* TEMP_TYPE)(INS);
		ins_desc_t * p = get_ins_desc(xed_icls);
		if( !p ) return -1;
		p->pre = (TEMP_TYPE)inspector;
		return 0;
	}
	void	mark_mem_taint	( size_t addr, size_t size, bool mark )
	{
		if( mark )
			return tagmap_setn( addr, size );
		else
			return tagmap_clrn( addr, size );
	}
	size_t	test_byte_taint	( size_t addr )
	{
		return tagmap_getb(addr);
	}
	size_t	test_word_taint	( size_t addr )
	{
		return tagmap_getw(addr);
	}
	size_t	test_dword_taint( size_t addr )
	{
		return tagmap_getl(addr);
	}
	size_t	test_qword_taint( size_t addr )
	{
		return tagmap_getq(addr);
	}
	size_t	test_mem_taint	( size_t addr, size_t size )
	{
		return tagmap_issetn(addr,size);
	}
	void	destroy()
	{
		libdft_die();
	}

	//////////////////////////////////////////////////////////////////////////

	void __Before_NtReadVirtualMemory( 
		ADDRINT hFile,
		ADDRINT lpBuffer,
		ADDRINT nNumberOfBytesToRead,
		ADDRINT lpNumberOfBytesRead )
	{
		if( hFile != COMM_DATA::MARK ) return ;

		if( nNumberOfBytesToRead < sizeof(COMM_DATA) ) 
			return ;

		COMM_DATA * xdata = (COMM_DATA * )lpBuffer;
		if( xdata->mark != COMM_DATA::MARK ) 
			return ;

		if( xdata->cmdcls == CMD_MARK_TAINT )
		{
			if( xdata->u.mark_taint.tof )
				tagmap_setn( xdata->u.mark_taint.buffer, xdata->u.mark_taint.bytes );
			else
				tagmap_clrn( xdata->u.mark_taint.buffer, xdata->u.mark_taint.bytes );
			xdata->lr = 0;
		} 
		else if( xdata->cmdcls == CMD_QUERY_TAINT )
		{
			size_t m = tagmap_issetn( xdata->u.query_taint.buffer, xdata->u.query_taint.bytes );
			if( m ) xdata->lr = 0;
		}
		else
		{
			xdata->lr = 0x8000FFFFL;
		}
		return ;
	}

	VOID Image(IMG img, VOID *v)
	{
		string name = IMG_Name(img);
		RTN readFileRtn = RTN_FindByName( img, "NtReadVirtualMemory" );
		if (!RTN_Valid(readFileRtn)) return ;
		RTN_Open(readFileRtn);
		RTN_InsertCall(readFileRtn, IPOINT_BEFORE, (AFUNPTR)__Before_NtReadVirtualMemory,
			IARG_FUNCARG_ENTRYPOINT_VALUE, 0,	// HANDLE
			IARG_FUNCARG_ENTRYPOINT_VALUE, 2,	// Buffer
			IARG_FUNCARG_ENTRYPOINT_VALUE, 3,	// SIZE
			IARG_FUNCARG_ENTRYPOINT_VALUE, 4,	// OUT SIZE
			IARG_END);
		RTN_Close(readFileRtn);
	}

	void dta_startup()
	{
		set_ins_inspect(XED_ICLASS_RET_NEAR, rop_check::trace_ret );
		IMG_AddInstrumentFunction(Image, 0);
	}


};
